/*
 * This file is part of FFmpeg.
 *
 * FFmpeg is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * FFmpeg is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */

#define _GNU_SOURCE
#include <fcntl.h>
#include <sys/mman.h>
#include <unistd.h>

#include <drm_fourcc.h>
#include <xf86drm.h>

#include "avassert.h"
#include "hwcontext.h"
#include "hwcontext_drm.h"
#include "hwcontext_internal.h"

/**
 * Copy from libdrm_macros.h while is not exposed by libdrm,
 * be replaced by #include "libdrm_macros.h" someday.
 */

/**
 * Static (compile-time) assertion.
 * Basically, use COND to dimension an array.  If COND is false/zero the
 * array size will be -1 and we'll get a compilation error.
 */
#define STATIC_ASSERT(COND) \
   do { \
      (void) sizeof(char [1 - 2*!(COND)]); \
   } while (0)

#if defined(ANDROID) && !defined(__LP64__)
#include <errno.h> /* for EINVAL */

extern void *__mmap2(void *, size_t, int, int, int, size_t);

static inline void *drm_mmap(void *addr, size_t length, int prot, int flags,
                             int fd, loff_t offset)
{
   /* offset must be aligned to 4096 (not necessarily the page size) */
   if (offset & 4095) {
      errno = EINVAL;
      return MAP_FAILED;
   }

   return __mmap2(addr, length, prot, flags, fd, (size_t) (offset >> 12));
}

#  define drm_munmap(addr, length) \
              munmap(addr, length)

#else

/* assume large file support exists */
#  define drm_mmap(addr, length, prot, flags, fd, offset) \
              mmap(addr, length, prot, flags, fd, offset)

static inline int drm_munmap(void *addr, size_t length)
{
   /* Copied from configure code generated by AC_SYS_LARGEFILE */
#define LARGE_OFF_T ((((off_t) 1 << 31) << 31) - 1 + \
                     (((off_t) 1 << 31) << 31))
   STATIC_ASSERT(LARGE_OFF_T % 2147483629 == 721 &&
                 LARGE_OFF_T % 2147483647 == 1);
#undef LARGE_OFF_T

   return munmap(addr, length);
}
#endif

// default
static int card_index = 0;

static const struct {
    enum AVPixelFormat pixfmt;
    uint32_t drm_format;
} supported_formats[] = {
    { AV_PIX_FMT_NV12,      DRM_FORMAT_NV12,        },
#ifdef DRM_FORMAT_NV12_10
    { AV_PIX_FMT_P010LE,    DRM_FORMAT_NV12_10,     },
    { AV_PIX_FMT_P010BE,    DRM_FORMAT_NV12_10  | DRM_FORMAT_BIG_ENDIAN },
#endif
    { AV_PIX_FMT_NV21,      DRM_FORMAT_NV21,        },
    { AV_PIX_FMT_YUV420P,   DRM_FORMAT_YUV420,      },
    { AV_PIX_FMT_YUYV422,   DRM_FORMAT_YUYV,        },
    { AV_PIX_FMT_YVYU422,   DRM_FORMAT_YVYU,        },
    { AV_PIX_FMT_UYVY422,   DRM_FORMAT_UYVY,        },
    { AV_PIX_FMT_NV16,      DRM_FORMAT_NV16,        },
    { AV_PIX_FMT_YUV422P,   DRM_FORMAT_YUV422,      },
#ifdef DRM_FORMAT_R16
    { AV_PIX_FMT_GRAY16LE,  DRM_FORMAT_R16,         },
    { AV_PIX_FMT_GRAY16BE,  DRM_FORMAT_R16      | DRM_FORMAT_BIG_ENDIAN },
#endif
    { AV_PIX_FMT_BGR8,      DRM_FORMAT_BGR233,      },
    { AV_PIX_FMT_RGB555LE,  DRM_FORMAT_XRGB1555,    },
    { AV_PIX_FMT_RGB555BE,  DRM_FORMAT_XRGB1555 | DRM_FORMAT_BIG_ENDIAN },
    { AV_PIX_FMT_BGR555LE,  DRM_FORMAT_XBGR1555,    },
    { AV_PIX_FMT_BGR555BE,  DRM_FORMAT_XBGR1555 | DRM_FORMAT_BIG_ENDIAN },
    { AV_PIX_FMT_RGB565LE,  DRM_FORMAT_RGB565,      },
    { AV_PIX_FMT_RGB565BE,  DRM_FORMAT_RGB565   | DRM_FORMAT_BIG_ENDIAN },
    { AV_PIX_FMT_BGR565LE,  DRM_FORMAT_BGR565,      },
    { AV_PIX_FMT_BGR565BE,  DRM_FORMAT_BGR565   | DRM_FORMAT_BIG_ENDIAN },
    { AV_PIX_FMT_RGB24,     DRM_FORMAT_RGB888,      },
    { AV_PIX_FMT_BGR24,     DRM_FORMAT_BGR888,      },
    { AV_PIX_FMT_0RGB,      DRM_FORMAT_BGRX8888,    },
    { AV_PIX_FMT_0BGR,      DRM_FORMAT_RGBX8888,    },
    { AV_PIX_FMT_RGB0,      DRM_FORMAT_XBGR8888,    },
    { AV_PIX_FMT_BGR0,      DRM_FORMAT_XRGB8888,    },
    { AV_PIX_FMT_ARGB,      DRM_FORMAT_BGRA8888,    },
    { AV_PIX_FMT_ABGR,      DRM_FORMAT_RGBA8888,    },
    { AV_PIX_FMT_RGBA,      DRM_FORMAT_ABGR8888,    },
    { AV_PIX_FMT_BGRA,      DRM_FORMAT_ARGB8888,    },
};

enum AVPixelFormat av_drm_get_pixfmt(uint32_t drm_format) {
    for (int i = 0; i < FF_ARRAY_ELEMS(supported_formats); i++) {
        if (supported_formats[i].drm_format == drm_format)
            return supported_formats[i].pixfmt;
    }
    return AV_PIX_FMT_NONE;
}

static void drm_device_free(AVHWDeviceContext *hwdev)
{
    AVDRMDeviceContext *hwctx = hwdev->hwctx;

    close(hwctx->fd);
}

static int drm_device_create(AVHWDeviceContext *hwdev, const char *device,
                             AVDictionary *opts, int flags)
{
    AVDRMDeviceContext *hwctx = hwdev->hwctx;
    drmVersionPtr version;
    char drm_dev[] = "/dev/dri/card0000";
    uint64_t has_dumb;

    if (!device) {
        snprintf(drm_dev, sizeof(drm_dev), DRM_DEV_NAME, DRM_DIR_NAME,
                 card_index);
        device = drm_dev;
    }
    hwctx->fd = open(device, O_RDWR);
    if (hwctx->fd < 0)
        return AVERROR(errno);

    version = drmGetVersion(hwctx->fd);
    if (!version) {
        av_log(hwdev, AV_LOG_ERROR, "Failed to get version information "
               "from %s: probably not a DRM device?\n", device);
        close(hwctx->fd);
        return AVERROR(EINVAL);
    }

    av_log(hwdev, AV_LOG_VERBOSE, "Opened DRM device %s: driver %s "
           "version %d.%d.%d.\n", device, version->name,
           version->version_major, version->version_minor,
           version->version_patchlevel);

    drmFreeVersion(version);

    if (drmGetCap(hwctx->fd, DRM_CAP_DUMB_BUFFER, &has_dumb) < 0 ||
        !has_dumb) {
        av_log(hwdev, AV_LOG_ERROR, "drm device '%s' "
               "does not support dumb buffers\n", device);
        close(hwctx->fd);
        return AVERROR(EINVAL);
    }

    hwdev->free = &drm_device_free;

    return 0;
}

static int drm_frames_get_constraints(AVHWDeviceContext *hwdev,
                                      const void *hwconfig,
                                      AVHWFramesConstraints *constraints)
{
    int i;

    constraints->min_width = 16;
    constraints->min_height = 16;

    constraints->valid_hw_formats =
        av_malloc_array(2, sizeof(enum AVPixelFormat));
    if (!constraints->valid_hw_formats)
        return AVERROR(ENOMEM);
    constraints->valid_hw_formats[0] = AV_PIX_FMT_DRM_PRIME;
    constraints->valid_hw_formats[1] = AV_PIX_FMT_NONE;

    constraints->valid_sw_formats =
        av_malloc_array(FF_ARRAY_ELEMS(supported_formats) + 1,
                        sizeof(enum AVPixelFormat));
    if (!constraints->valid_sw_formats)
        return AVERROR(ENOMEM);
    for(i = 0; i < FF_ARRAY_ELEMS(supported_formats); i++)
        constraints->valid_sw_formats[i] = supported_formats[i].pixfmt;
    constraints->valid_sw_formats[i] = AV_PIX_FMT_NONE;

    return 0;
}

static void free_drm_frame_descriptor(AVDRMDeviceContext *hwctx,
                                      AVDRMFrameDescriptor *desc)
{
    int i;
    if (!desc)
        return;

    for (i = 0; i < desc->nb_objects; i++) {
        AVDRMObjectDescriptor *object = &desc->objects[i];
        if (object->ptr)
            drm_munmap(object->ptr, object->size);
        if (object->fd > 0) {
            int ret;
            uint32_t handle = 0;

            ret = drmPrimeFDToHandle(hwctx->fd, object->fd, &handle);
            if (ret)
                av_log(NULL, AV_LOG_WARNING,
                       "Failed to convert drm fd to handle: %m\n");
            if (handle > 0) {
                struct drm_mode_destroy_dumb data = {
                    .handle = handle,
                };
                ret = drmIoctl(hwctx->fd, DRM_IOCTL_MODE_DESTROY_DUMB, &data);
                if (ret)
                    av_log(NULL, AV_LOG_WARNING,
                           "Failed to free drm handle: %m\n");
            }

            ret = close(object->fd);
            if (ret)
                av_log(NULL, AV_LOG_WARNING,
                       "Failed to close drm buffer fd = %d: %m\n", object->fd);
        }
    }

    memset(desc, 0, sizeof(*desc));
    av_free(desc);
}

static void drm_buffer_free(void *opaque, uint8_t *data)
{
    AVHWFramesContext *hwfc = opaque;
    AVDRMDeviceContext *hwctx = hwfc->device_ctx->hwctx;
    AVDRMFrameDescriptor *desc = (AVDRMFrameDescriptor *)data;

    free_drm_frame_descriptor(hwctx, desc);
}

static AVBufferRef *drm_pool_alloc(void *opaque, int size)
{
    int ret;
    AVHWFramesContext *hwfc = opaque;
    AVDRMDeviceContext *hwctx = hwfc->device_ctx->hwctx;
    AVDRMFrameDescriptor *desc;
    AVDRMLayerDescriptor *layer;
    AVBufferRef *ref;

    int i;
    const AVPixFmtDescriptor *pixdesc = av_pix_fmt_desc_get(hwfc->sw_format);
    struct drm_mode_create_dumb dmcb;
    struct drm_mode_map_dumb dmmd;

    desc = av_mallocz(sizeof(*desc));
    if (!desc)
        return NULL;

    memset(&dmcb, 0, sizeof(struct drm_mode_create_dumb));
    dmcb.bpp = av_get_bits_per_pixel(pixdesc);
    dmcb.width  = FFALIGN(hwfc->width, 16);
    dmcb.height = FFALIGN(hwfc->height, 16);
    ret = drmIoctl(hwctx->fd, DRM_IOCTL_MODE_CREATE_DUMB, &dmcb);
    if (ret < 0) {
        av_log(hwfc, AV_LOG_ERROR,
               "Failed to create dumb<w,h,bpp: %d,%d,%d>: %m.\n",
               dmcb.width, dmcb.height, dmcb.bpp);
        goto fail;
    }
    av_assert0(dmcb.size >= dmcb.width * dmcb.height * dmcb.bpp / 8);

    desc->nb_objects = 1;
    desc->nb_layers = 1;
    ret = drmPrimeHandleToFD(hwctx->fd, dmcb.handle, DRM_CLOEXEC,
                             &desc->objects[0].fd);
    if (ret) {
        av_log(hwfc, AV_LOG_ERROR, "Failed to convert handle to fd: %m\n");
        goto fail;
    }
    memset(&dmmd, 0, sizeof(dmmd));
    dmmd.handle = dmcb.handle;

    ret = drmIoctl(hwctx->fd, DRM_IOCTL_MODE_MAP_DUMB, &dmmd);
    if (ret) {
        av_log(hwfc, AV_LOG_ERROR, "Failed to map dumb: %m\n");
        goto fail;
    }

    // default read and write
    desc->objects[0].ptr = drm_mmap(NULL, dmcb.size, PROT_READ | PROT_WRITE,
                                    MAP_SHARED, hwctx->fd, dmmd.offset);
    if (desc->objects[0].ptr == MAP_FAILED) {
        av_log(hwfc, AV_LOG_ERROR, "Failed to drm_mmap: %m\n");
        goto fail;
    }

    desc->objects[0].size = dmcb.size;

    layer = &desc->layers[0];
    for (i = 0; i < FF_ARRAY_ELEMS(supported_formats); i++) {
        if (supported_formats[i].pixfmt == hwfc->sw_format) {
            layer->format = supported_formats[i].drm_format;
            break;
        }
    }
    layer->nb_planes = av_pix_fmt_count_planes(hwfc->sw_format);
    layer->planes[0].object_index = 0;
    layer->planes[0].offset = 0;
    layer->planes[0].pitch =
        av_image_get_linesize(hwfc->sw_format, hwfc->width, 0);

    for (i = 1; i < layer->nb_planes; i++) {
        layer->planes[i].object_index = 0;
        layer->planes[i].offset = layer->planes[i-1].pitch * hwfc->height;
        layer->planes[i].pitch =
            av_image_get_linesize(hwfc->sw_format, hwfc->width, i);
    }

    ref = av_buffer_create((uint8_t*)desc, sizeof(*desc), drm_buffer_free,
                           opaque, 0);
    if (!ref) {
        av_log(hwfc, AV_LOG_ERROR, "Failed to create drm buffer.\n");
        goto fail;
    }

    return ref;

fail:
    free_drm_frame_descriptor(hwctx, desc);
    return NULL;
}

static int drm_frames_init(AVHWFramesContext *hwfc)
{
    int i;
    if (hwfc->pool) {
        // has been set outside?
        return 0;
    }

    for (i = 0; i < FF_ARRAY_ELEMS(supported_formats); i++)
        if (supported_formats[i].pixfmt == hwfc->sw_format)
            break;
    if (i >= FF_ARRAY_ELEMS(supported_formats)) {
        av_log(hwfc, AV_LOG_ERROR, "Unsupported format: %s.\n",
               av_get_pix_fmt_name(hwfc->sw_format));
        return AVERROR(EINVAL);
    }

    hwfc->internal->pool_internal =
        av_buffer_pool_init2(sizeof(AVDRMFrameDescriptor), hwfc, drm_pool_alloc,
                             NULL);
    if (!hwfc->internal->pool_internal) {
        av_log(hwfc, AV_LOG_ERROR, "Failed to create drm buffer pool.\n");
        return AVERROR(ENOMEM);
    }

    return 0;
}

static void drm_frames_uninit(AVHWFramesContext *hwfc av_unused)
{}

static int drm_get_buffer(AVHWFramesContext *hwfc, AVFrame *frame)
{
    frame->buf[0] = av_buffer_pool_get(hwfc->pool);
    if (!frame->buf[0])
        return AVERROR(ENOMEM);

    frame->data[0] = (uint8_t*)frame->buf[0]->data;

    frame->format = AV_PIX_FMT_DRM_PRIME;
    frame->width  = hwfc->width;
    frame->height = hwfc->height;

    return 0;
}

typedef struct DRMMapping {
    // Address and length of each mmap()ed region.
    int nb_regions;
    void *address[AV_DRM_MAX_PLANES];
    size_t length[AV_DRM_MAX_PLANES];
} DRMMapping;

static void drm_unmap_frame(AVHWFramesContext *hwfc,
                            HWMapDescriptor *hwmap)
{
    DRMMapping *map = hwmap->priv;
    int i;

    for (i = 0; i < map->nb_regions; i++)
        munmap(map->address[i], map->length[i]);

    av_free(map);
}

static int drm_map_frame(AVHWFramesContext *hwfc,
                         AVFrame *dst, const AVFrame *src, int flags)
{
    const AVDRMFrameDescriptor *desc = (AVDRMFrameDescriptor*)src->data[0];
    DRMMapping *map;
    int err, i, p, plane;
    int mmap_prot;
    void *addr;

    map = av_mallocz(sizeof(*map));
    if (!map)
        return AVERROR(ENOMEM);

    mmap_prot = 0;
    if (flags & AV_HWFRAME_MAP_READ)
        mmap_prot |= PROT_READ;
    if (flags & AV_HWFRAME_MAP_WRITE)
        mmap_prot |= PROT_WRITE;

    av_assert0(desc->nb_objects <= AV_DRM_MAX_PLANES);
    for (i = 0; i < desc->nb_objects; i++) {
        addr = mmap(NULL, desc->objects[i].size, mmap_prot, MAP_SHARED,
                    desc->objects[i].fd, 0);
        if (addr == MAP_FAILED) {
            err = AVERROR(errno);
            av_log(hwfc, AV_LOG_ERROR, "Failed to map DRM object %d to "
                   "memory: %d.\n", desc->objects[i].fd, errno);
            goto fail;
        }

        map->address[i] = addr;
        map->length[i]  = desc->objects[i].size;
    }
    map->nb_regions = i;

    plane = 0;
    for (i = 0; i < desc->nb_layers; i++) {
        const AVDRMLayerDescriptor *layer = &desc->layers[i];
        for (p = 0; p < layer->nb_planes; p++) {
            dst->data[plane] =
                (uint8_t*)map->address[layer->planes[p].object_index] +
                                       layer->planes[p].offset;
            dst->linesize[plane] =     layer->planes[p].pitch;
            ++plane;
        }
    }
    av_assert0(plane <= AV_DRM_MAX_PLANES);

    dst->width  = src->width;
    dst->height = src->height;

    err = ff_hwframe_map_create(src->hw_frames_ctx, dst, src,
                                &drm_unmap_frame, map);
    if (err < 0)
        goto fail;

    return 0;

fail:
    for (i = 0; i < desc->nb_objects; i++) {
        if (map->address[i])
            munmap(map->address[i], map->length[i]);
    }
    av_free(map);
    return err;
}

static int drm_transfer_get_formats(AVHWFramesContext *ctx,
                                    enum AVHWFrameTransferDirection dir,
                                    enum AVPixelFormat **formats)
{
    enum AVPixelFormat *pix_fmts;

    pix_fmts = av_malloc_array(2, sizeof(*pix_fmts));
    if (!pix_fmts)
        return AVERROR(ENOMEM);

    pix_fmts[0] = ctx->sw_format;
    pix_fmts[1] = AV_PIX_FMT_NONE;

    *formats = pix_fmts;
    return 0;
}

static int drm_transfer_data_from(AVHWFramesContext *hwfc,
                                  AVFrame *dst, const AVFrame *src)
{
    AVFrame *map;
    int err;

    if (dst->width > hwfc->width || dst->height > hwfc->height)
        return AVERROR(EINVAL);

    map = av_frame_alloc();
    if (!map)
        return AVERROR(ENOMEM);
    map->format = dst->format;

    err = drm_map_frame(hwfc, map, src, AV_HWFRAME_MAP_READ);
    if (err)
        goto fail;

    map->width  = dst->width;
    map->height = dst->height;

    err = av_frame_copy(dst, map);
    if (err)
        goto fail;

    err = 0;
fail:
    av_frame_free(&map);
    return err;
}

static int drm_transfer_data_to(AVHWFramesContext *hwfc,
                                AVFrame *dst, const AVFrame *src)
{
    AVFrame *map;
    int err;

    if (src->width > hwfc->width || src->height > hwfc->height)
        return AVERROR(EINVAL);

    map = av_frame_alloc();
    if (!map)
        return AVERROR(ENOMEM);
    map->format = src->format;

    err = drm_map_frame(hwfc, map, dst, AV_HWFRAME_MAP_WRITE |
                                        AV_HWFRAME_MAP_OVERWRITE);
    if (err)
        goto fail;

    map->width  = src->width;
    map->height = src->height;

    err = av_frame_copy(map, src);
    if (err)
        goto fail;

    err = 0;
fail:
    av_frame_free(&map);
    return err;
}

static int drm_map_from(AVHWFramesContext *hwfc, AVFrame *dst,
                        const AVFrame *src, int flags)
{
    int err;

    if (hwfc->sw_format != dst->format)
        return AVERROR(ENOSYS);

    err = drm_map_frame(hwfc, dst, src, flags);
    if (err)
        return err;

    err = av_frame_copy_props(dst, src);
    if (err)
        return err;

    return 0;
}

const HWContextType ff_hwcontext_type_drm = {
    .type                   = AV_HWDEVICE_TYPE_DRM,
    .name                   = "DRM",

    .device_hwctx_size      = sizeof(AVDRMDeviceContext),

    .device_create          = &drm_device_create,

    .frames_get_constraints = &drm_frames_get_constraints,

    .frames_get_buffer      = &drm_get_buffer,
    .frames_init            = &drm_frames_init,
    .frames_uninit          = &drm_frames_uninit,
    .transfer_get_formats   = &drm_transfer_get_formats,
    .transfer_data_to       = &drm_transfer_data_to,
    .transfer_data_from     = &drm_transfer_data_from,
    .map_from               = &drm_map_from,

    .pix_fmts = (const enum AVPixelFormat[]) {
        AV_PIX_FMT_DRM_PRIME,
        AV_PIX_FMT_NONE
    },
};
